{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
    "\n",
    "from tldextract import TLDExtract\n",
    "extract = TLDExtract(suffix_list_urls=None)\n",
    "\n",
    "import time\n",
    "from numpy import mat\n",
    "from kafka import KafkaConsumer\n",
    "from keras.models import load_model\n",
    "\n",
    "\n",
    "domain2int={\"a\":1, \"b\":2, \"c\":3, \"d\":4, \"e\":5, \"f\":6, \n",
    "            \"g\":7, \"h\":8, \"i\":9, \"j\":10, \"k\":11, \"l\":12, \n",
    "            \"m\":13, \"n\":14, \"o\":15, \"p\":16, \"q\":17, \"r\":18, \n",
    "            \"s\":19, \"t\":20, \"u\":21, \"v\":22, \"w\":23, \"x\":24, \n",
    "            \"y\":25, \"z\":26, \"1\":27, \"2\":28, \"3\":29, \"4\":30, \n",
    "            \"5\":31, \"6\":32, \"7\":33, \"8\":34, \"9\":35, \"0\":36, \n",
    "            \"-\":37}\n",
    "\n",
    "def pad(dat, length=31, item=0):\n",
    "    if len(dat) <= 31:\n",
    "        dat.extend((length - len(dat))*[item])\n",
    "    else:\n",
    "        dat = dat[:length]\n",
    "    return dat\n",
    "\n",
    "def domain2list(domain):\n",
    "    data = list(map(lambda x:domain2int.get(x, 38), domain))\n",
    "    return pad(data)\n",
    "\n",
    "\n",
    "results = (\"clean\",\"DGA-domain\")\n",
    "\n",
    "false = False\n",
    "true = True\n",
    "\n",
    "model_fname=\"./models/BiLST-20180821-110303.module\"\n",
    "model = load_model(model_fname)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "clean www.baidu.com\n",
      "clean www.baidu.com\n",
      "clean ss1.bdstatic.com\n",
      "clean s1.bdstatic.com\n",
      "clean ss1.bdstatic.com\n",
      "clean ss1.bdstatic.com\n",
      "clean s1.bdstatic.com\n",
      "clean s1.bdstatic.com\n",
      "clean b1.bdstatic.com\n",
      "clean t1.baidu.com\n",
      "clean b1.bdstatic.com\n",
      "clean t1.baidu.com\n",
      "clean t11.baidu.com\n",
      "clean b1.bdstatic.com\n",
      "clean t12.baidu.com\n",
      "clean b1.bdstatic.com\n",
      "clean t11.baidu.com\n",
      "clean t2.baidu.com\n",
      "clean t12.baidu.com\n",
      "clean t2.baidu.com\n",
      "clean t3.baidu.com\n",
      "clean ss1.bdstatic.com\n",
      "clean b1.bdstatic.com\n",
      "clean t10.baidu.com\n",
      "clean t10.baidu.com\n",
      "clean t10.baidu.com\n",
      "clean t3.baidu.com\n",
      "clean s1.bdstatic.com\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean safebrowsing.googleapis.com\n",
      "clean safebrowsing.googleapis.com\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean ssl.gstatic.com\n",
      "clean www.gstatic.com\n",
      "clean ssl.gstatic.com\n",
      "clean www.gstatic.com\n",
      "clean ssl.gstatic.com\n",
      "clean www.gstatic.com\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n",
      "clean fedoraproject.org\n"
     ]
    }
   ],
   "source": [
    "consumer = KafkaConsumer('dns')\n",
    "for message in consumer:\n",
    "    # message value and key are raw bytes -- decode if necessary!\n",
    "    # e.g., for unicode: `message.value.decode('utf-8')`\n",
    "    dns_log = eval(message.value.decode('utf-8'))\n",
    "    dns_query = dns_log[\"query\"]\n",
    "    x = domain2list(extract(dns_query).domain)\n",
    "    X = mat([x,])\n",
    "    rs = results[model.predict_classes(X)[0][0]]\n",
    "    \n",
    "    print(rs, dns_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
